#GPU训练
import torch
from torch import nn
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader
# import torch

#1.准备训练数据 测试数据
train_data = torchvision.datasets.CIFAR10(
    root="./",
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=True
)
test_data = torchvision.datasets.CIFAR10(
    root="./",
    train=False,
    transform=torchvision.transforms.ToTensor(),
    download=True
)
# Ftrain_data = torchvision.datasets.coco
train_size = len(train_data)
test_size = len(test_data)
#加载数据集
train_dataloader = DataLoader(train_data, batch_size=128, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=128, shuffle=False)

#2.创建神经网络
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
            
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 64, 5, 1, 2),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Flatten(),
            nn.Linear(64*4*4, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        x = self.model(x)
        return x
    
#3.实例化模型
tudui = Tudui()
if torch.cuda.is_available():
    tudui = tudui.cuda()
    print("GPU训练")
else:
    print("CPU训练")

#4.损失函数与优化器
loss_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():
    loss_fn = loss_fn.cuda()
learning_rate = 1e-3
optimizer = torch.optim.Adam(tudui.parameters(), lr= learning_rate)

#5.训练与测试
epoch = 10
for i in range(epoch):
    acy = 0
    total_loss = 0
    tudui.train()
    for data in train_dataloader:
        imgs, targets = data
        if torch.cuda.is_available():
            imgs, targets = imgs.cuda(), targets.cuda()
        outputs = tudui(imgs)
        acy_t = (outputs.argmax(1) == targets).sum()
        acy += acy_t
        loss = loss_fn(outputs, targets)
        # print("this batch of loss:{}".format(loss.item()))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss+=loss.item()
    print("训练epoch:{}, loss:{} acy:{}".format(i+1, total_loss, acy/train_size))


#6.测试
accuracy = 0
tudui.eval()
with torch.no_grad():
    for data in test_dataloader:
        imgs, targets = data
        if torch.cuda.is_available():
            imgs, targets = imgs.cuda(), targets.cuda()
        outputs = tudui(imgs)
        accuracy_t = (outputs.argmax(1) == targets).sum()
        accuracy += accuracy_t
print("测试集合的正确率为:{}".format(accuracy/test_size))







































